package com.xiashitech.agent.instrumentation.spring.controller;

import cn.hutool.http.HttpResponse;
import cn.hutool.http.HttpUtil;
import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.db.entity.DataWaterMarkConfig;
import com.xiashitech.agent.dto.UserRequest;
import com.xiashitech.agent.instrumentation.base.AbstractAttributesExtractor;
import com.xiashitech.agent.instrumentation.base.context.TraceContextVars;
import com.xiashitech.agent.instrumentation.base.context.XSBaggage;
import com.xiashitech.agent.instrumentation.base.context.bridge.BridgeBaggageBuilder;
import com.xiashitech.agent.instrumentation.biz.api.increment.ApiIncrementHelper;
import com.xiashitech.agent.instrumentation.bridge.Java8Bridge;
import com.xiashitech.agent.instrumentation.httprequest.XSHttpServletResponseWrapper;
import com.xiashitech.agent.instrumentation.spring.controller.apicapture.cache.ApiCache;
import com.xiashitech.agent.instrumentation.spring.controller.apicapture.inst.utils.GenerateSpanAttrUtil;
import com.xiashitech.agent.instrumentation.spring.controller.apicapture.inst.v_3_1.ApiGenerate_V_3_1;
import com.xiashitech.agent.instrumentation.spring.controller.apicapture.utils.ApiSyncUtil;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import com.xiashitech.agent.instrumentation.utils.MachineAttributesUtil;
import com.xiashitech.agent.proxy.hutool.StrUtil;
import com.xiashitech.agent.service.extload.AuthQueryService;
import com.xiashitech.agent.utils.JsonUtil;
import com.xiashitech.agent.utils.XSUtil;
import com.xiashitech.interfaces.agent.dto.AgentConfigDTO;
import com.xiashitech.interfaces.agent.dto.SystemInfoDTO;
import com.xiashitech.interfaces.agent.dto.api.Api;
import com.xiashitech.interfaces.agent.dto.api.DesenseParam;
import io.opentelemetry.javaagent.bootstrap.internal.InstrumentationConfig;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.baggage.Baggage;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.common.AttributesBuilder;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.trace.Span;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.Nullable;

import javax.json.*;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.CharArrayWriter;
import java.io.IOException;
import java.io.StringReader;
import java.lang.reflect.Method;
import java.util.*;

public class XSControllerRequestAttributesExtractor extends AbstractAttributesExtractor<UserRequest, Void> {
    public String getAttributesExtractorName() {
        return XSControllerRequestAttributesExtractor.class.getName();
    }

    @Override
    public void onXSStart(AttributesBuilder attributesBuilder, Context context, UserRequest userRequest) {
        Baggage baggage = Baggage.fromContext(context);
        System.out.println("###spring onXSStart TraceId "+ Span.current().getSpanContext().getTraceId());
        try {
            String xsBaggageUserId = Java8Bridge.getUserIdFromXSBaggage();
            System.out.println("###spring1 xsBaggageUserId "+xsBaggageUserId);
            String baggageUserId = baggage.getEntryValue("xs.userId");
            if(baggageUserId == null || baggageUserId.isEmpty()) {
                baggage = baggage.toBuilder().put("xs.userId",xsBaggageUserId).build();
                TraceContextVars contextVars = TraceContextVars.get();
                baggage.forEach((key,value) -> contextVars.set(key,value.getValue()));
            }
            if ((userRequest.getUserid() == null || "".equals(userRequest.getUserid())) && baggage.size() > 0) {
                userRequest.setUserid(baggage.getEntryValue("xs.userId"));
                userRequest.setRequestDate(baggage.getEntryValue("xs.requestDate"));
                userRequest.setSessionId(baggage.getEntryValue("xs.sessionId"));
                userRequest.setSessionUpdate(baggage.getEntryValue("xs.sessionUpdate") == null ? false : Boolean.valueOf(baggage.getEntryValue("xs.sessionUpdate")));
            }
            if (baggage.size() > 0 && baggage.getEntryValue("xs.loginName") != null) {
                userRequest.setLoginName(baggage.getEntryValue("xs.loginName"));
            }
            System.out.println("###spring AgentConfigDTO.getConfig().getCaptureApi() "+AgentConfigDTO.getConfig().getCaptureApi());
            System.out.println("###spring8 AgentConfigDTO.getConfig() "+AgentConfigDTO.getConfig());
            if (AgentConfigDTO.getConfig().getCaptureApi() != null && AgentConfigDTO.getConfig().getCaptureApi()) {
                //生成api接口信息
                System.out.println("###spring9 userRequest.getAdviceObject().getClass().getName() "+userRequest.getAdviceObject().getClass().getName());

                if (userRequest.getAdviceObject().getClass().getName().equals("org.springframework.web.servlet.mvc.method.annotation.ServletInvocableHandlerMethod")) {
                    System.out.println("###spring10");
                    Object servletInvocableHandlerMethod = userRequest.getAdviceObject();
                    GenerateSpanAttrUtil.generateSpanAttr(userRequest, attributesBuilder, (Class) servletInvocableHandlerMethod.getClass().getMethod("getBeanType").invoke(servletInvocableHandlerMethod), (Method) servletInvocableHandlerMethod.getClass().getMethod("getMethod").invoke(servletInvocableHandlerMethod));
                    System.out.println("###spring11");
                    /**
                     * 将api名字存入tracing返回上下文
                     */
                    String endpointApiIsNull = baggage.getEntryValue("endpoint-api");
                    System.out.println(Thread.currentThread().getId() + " DesenseDebug, XSControllerRequestAttributesExtractor onXSStart endpointApiIsNull: " + endpointApiIsNull);
                    if (endpointApiIsNull == null || endpointApiIsNull.isEmpty()) {
                        String methodSign = GenerateSpanAttrUtil.getMethodSign((Class) servletInvocableHandlerMethod.getClass().getMethod("getBeanType").invoke(servletInvocableHandlerMethod), (Method) servletInvocableHandlerMethod.getClass().getMethod("getMethod").invoke(servletInvocableHandlerMethod));
                        Api api = ApiCache.cache.get(methodSign);
                        StringJoiner stringJoiner = new StringJoiner(",");
                        api.getPaths().forEach(path -> stringJoiner.add(path.getName()));
                        String endpointApi = stringJoiner.toString();
                        System.out.println(Thread.currentThread().getId() + " DesenseDebug, XSControllerRequestAttributesExtractor onXSStart endpointApi: " + endpointApi);
                        userRequest.setEndpointApi(endpointApi);
                        userRequest.setApiSystemId(StartConfig.getParamValueFromKey(StartConfig.systemId));
                        userRequest.setIsOutApi(Boolean.TRUE);
                    }
                }
            }
        } catch (Throwable e) {
            e.printStackTrace();
            StartConfig.throwableStackPrinter(e);
        }
    }

    @Override
    public void onXSEnd(AttributesBuilder attributesBuilder, Context context, UserRequest userRequest, @Nullable Void unused, @Nullable Throwable throwable) {
        if (userRequest != null) {
            try {
                attributesBuilder.put("xs.userId", userRequest.getUserid() == null ? "" : userRequest.getUserid());
                attributesBuilder.put("xs.loginName", userRequest.getLoginName() == null ? "" : userRequest.getLoginName());
                attributesBuilder.put("xs.sessionId", userRequest.getSessionId() == null ? "" : userRequest.getSessionId());
                attributesBuilder.put("xs.requestDate", userRequest.getRequestDate() == null ? "" : userRequest.getRequestDate().toString());
                attributesBuilder.put("xs.sessionUpdate", userRequest.getSessionUpdate() == null ? false : userRequest.getSessionUpdate());
                attributesBuilder.put("xs.systemId", StartConfig.getParamValueFromKey(StartConfig.systemId));
                attributesBuilder.put("xs.systemReg", SystemInfoDTO.getConfig().getSystemExist());
                attributesBuilder.put("xs.sourceSystemId", BaggageUtil.getTraceVar(context, "xsAuthSystemId") == null ? StartConfig.getParamValueFromKey(StartConfig.systemId) : BaggageUtil.getTraceVar(context, "xsAuthSystemId"));
                TraceContextVars.get().set("xs.sessionUpdate", userRequest.getSessionUpdate() == null ? "false" : userRequest.getSessionUpdate().toString());
                Object servletWebRequest = userRequest.getRequest();
                Object response = servletWebRequest.getClass().getMethod("getResponse").invoke(servletWebRequest);
                String contentType = (String) response.getClass().getMethod("getContentType").invoke(response);
                // 泛化 url 的出参，采用这方法，支持类上或接口上注解的路径
                Api api = ApiGenerate_V_3_1.generateApi(userRequest);
                /*
                Object servletInvocableHandlerMethod = userRequest.getAdviceObject();
                Method controllerMethod = (Method) servletInvocableHandlerMethod.getClass().getMethod("getMethod").invoke(servletInvocableHandlerMethod);
                Annotation[] methodAnnotations = controllerMethod.getAnnotations();
                for(Annotation annotation: methodAnnotations) {
                    String typeName = annotation.annotationType().getName();
                    if(typeName.equals("org.springframework.web.bind.annotation.PutMapping")
                         || typeName.equals("org.springframework.web.bind.annotation.DeleteMapping")
                         || typeName.equals("org.springframework.web.bind.annotation.GetMapping")
                         || typeName.equals("org.springframework.web.bind.annotation.PatchMapping")
                         || typeName.equals("org.springframework.web.bind.annotation.PostMapping")
                         || typeName.equals("org.springframework.web.bind.annotation.RequestMapping")) {
                        String[] values = (String[]) annotation.getClass().getMethod("value").invoke(annotation);
                        String mapPath = "";
                        long ts = System.currentTimeMillis();
                        if(values.length == 0) {
                            String[] paths = (String[]) annotation.getClass().getMethod("path").invoke(annotation);
                            if(paths.length > 0)
                                mapPath = paths[0];
                            for(String str : paths) System.out.println(ts + ", XSControllerRequestAttributesExtractorA: " + str);
                        } else {
                            mapPath = values[0];
                            for(String str : values) System.out.println(ts + ", XSControllerRequestAttributesExtractorB: " + str);
                        }
                        if(mapPath != null && !mapPath.isEmpty()) {
                            System.out.println(ts + ", XSControllerRequestAttributesExtractorA: " + mapPath);
                            ApiIncrementHelper.syncIncrementApiByGenericType(servletWebRequest, response, (String) mapPath);
                        }
                    }
                }
                */
                if(api != null && api.getPaths() != null && api.getPaths().size() > 0) {
                    String mapPath = api.getPaths().get(0).getName(); // TODO 只取第一个
                    if (mapPath != null && !mapPath.isEmpty()) {
                        ApiIncrementHelper.syncIncrementApiByGenericType(servletWebRequest, response, (String) mapPath);
                    }
                }

                Boolean httpWrapper = null;
                if (AgentConfigDTO.getConfig().getHttpWrapper() != null) {
                    httpWrapper = AgentConfigDTO.getConfig().getHttpWrapper().get(StartConfig.getParamValueFromKey(StartConfig.systemId));
                }
                if ((httpWrapper == null || httpWrapper) && !(response instanceof XSHttpServletResponseWrapper) && response instanceof HttpServletResponseWrapper) {
                    boolean isXSHttpServletResponseWrapper = (boolean) response.getClass().getMethod("isWrapperFor", Class.class).invoke(response, XSHttpServletResponseWrapper.class);
                    if (isXSHttpServletResponseWrapper) {
                        while (true) {
                            Object innerResponse = response.getClass().getMethod("getResponse").invoke(response);
                            if (innerResponse instanceof XSHttpServletResponseWrapper) {
                                response = innerResponse;
                                break;
                            } else {
                                response = innerResponse;
                            }
                        }
                    }
                }
                String desenseImpl = null;
                Map<String, String> desenseImplMap = AgentConfigDTO.getConfig().getDesenseImpl();
                if (desenseImplMap != null){
                    desenseImpl = desenseImplMap.get(StartConfig.getParamValueFromKey(StartConfig.systemId));
                }
                String charset = InstrumentationConfig.get().getString("xs.charset.textcontent");
                if (desenseImpl == null || "".equals(desenseImpl) || "springWebmvc".equals(desenseImpl)) {
                    if (contentType != null && contentType.contains("json") && response instanceof XSHttpServletResponseWrapper) {
                        XSHttpServletResponseWrapper xsHttpServletResponseWrapper = (XSHttpServletResponseWrapper) response;
                        byte[] responseWrapperByteArray = xsHttpServletResponseWrapper.getByteArray();
                        char[] responseWrapperCharArray = xsHttpServletResponseWrapper.getCharArray();
                        /**
                         * 根据数据访问权限及配置项对api返回值进行脱敏
                         */
                        if (userRequest.getIsOutApi()) {
                            String jsonString = null;
                            if (responseWrapperByteArray.length > 0) {
                                if (charset != null) {
                                    jsonString = new String(responseWrapperByteArray, charset);
                                } else {
                                    jsonString = new String(responseWrapperByteArray);
                                }
                            } else if (responseWrapperCharArray.length > 0) {
                                jsonString = new String(responseWrapperCharArray);
                            }
                            /**
                             * 获取api的返回值的key列表
                             */
                            if (jsonString != null) {
                                Set<String> jsonKeySet = new HashSet();
                                JsonUtil.getJsonKeys(jsonString, jsonKeySet);
                                String endpointApi = userRequest.getEndpointApi();
                                if (endpointApi != null) {
                                    endpointApi = endpointApi.replace("/", "_");
                                }
                                DesenseParam desenseParam = DesenseParam.builder().apiName(endpointApi).system(userRequest.getApiSystemId()).desenseParam(new ArrayList(jsonKeySet)).build();
                                ApiSyncUtil.syncApiDesense(desenseParam);
                            }
                            XSBaggage[] baggagesWithCurrentAndConfluence = XSBaggage.currentWithConfluence();
                            for (XSBaggage xsBaggage : baggagesWithCurrentAndConfluence) {
                                if (xsBaggage != null) {
                                    String endpointApiBack = xsBaggage.getEntryValue("endpoint-api-back");
                                    if (jsonString != null && endpointApiBack != null && !endpointApiBack.isEmpty()) {
                                        String jsonPath = jsonString;
                                        String[] endpointApiBacks = endpointApiBack.split("#");
                                        for (String endpointApiBackItem : endpointApiBacks) {
                                            System.out.println(Thread.currentThread().getId() + " DesenseDebug, XSControllerRequestAttributesExtractor onXSEnd endpointApiBackItem: " + endpointApiBackItem);
                                            String apiName = null;
                                            String[] desenseAction = endpointApiBackItem.split("&");
                                            for (int i = 0; i < desenseAction.length; i++) {
                                                if (i == 0) {
                                                    apiName = desenseAction[0];
                                                    if (!userRequest.getEndpointApi().equals(apiName)) {
                                                        break;
                                                    } else {
                                                        continue;
                                                    }
                                                }
                                                String[] desenseActionItem = desenseAction[i].split("=");
                                                String patchPath = desenseActionItem[0];
                                                String[] patchPathItem = desenseActionItem[1].split(",");
                                                String isExecute = patchPathItem[0];
                                                String regular = patchPathItem[1];
                                                String replaceSymbol = patchPathItem[2];
                                                String systemId = patchPathItem[3];
                                                //执行脱敏
                                                if ("true".equals(isExecute) && StartConfig.getParamValueFromKey(StartConfig.systemId).equals(systemId)) {
                                                    jsonPath = JsonUtil.jsonPatch(jsonPath, patchPath, regular, replaceSymbol);
                                                    if (jsonPath != null) {
                                                        writeResponse(xsHttpServletResponseWrapper, jsonPath, charset);
                                                    }
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
                if (response instanceof XSHttpServletResponseWrapper && CacheStatus.isDwmEnable()) {
                    addDataWaterMark(userRequest, (XSHttpServletResponseWrapper) response, charset);
                }
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            finally {
                Java8Bridge.setUserIdToXSBaggage(null);
            }
        }
        MachineAttributesUtil.fillMachineAttributes(attributesBuilder);
    }

    /**
     * 根据水印配置给controller返回值增加水印
     * @param request
     * @param response
     * @return
     */
    private void addDataWaterMark(UserRequest request, XSHttpServletResponseWrapper response, String charset) {
        String apiPath = getApiPath(request);
        String userId = request.getUserid();
        String systemName = StartConfig.getParamValueFromKey(StartConfig.systemId);
        String waterMarkService = AgentConfigDTO.getConfig().getWaterMarkService();
        if (StringUtils.isAnyBlank(apiPath, systemName, userId, waterMarkService)) {
            XSUtil.warn(this.getClass(), "DWM missing parameter. apiPath:[{}], userId:[{}], waterMarkService:[{}]", apiPath, userId, waterMarkService);
        } else {
            String jsonBody = response.getJsonBody();
            if (!JsonUtil.isJson(jsonBody)) {
                return;
            }
            DataWaterMarkConfig config = AuthQueryService.getInstance().findDataWaterMarkConfig(systemName, apiPath);
            if (config == null) {
                return;
            }
            Integer dwmThreshold = config.getMorethanLine();
            if (dwmThreshold == null || dwmThreshold < 1) {
                return;
            }
            String url = XSUtil.makeHttpUrl(StrUtil.format(waterMarkService, config.getId(), systemName, userId));
            try {
                if (JsonUtil.isJsonArray(jsonBody)) {
                    if (JsonUtil.getJsonArraySize(jsonBody) < dwmThreshold) {
                        return;
                    }
                    jsonBody = getDataWaterMarkResult(url, jsonBody);
                    writeResponse(response, jsonBody, charset);
                }
                if (JsonUtil.isJsonObj(jsonBody)) {
                    jsonBody = addWaterMarkByPath(url, jsonBody, dwmThreshold);
                    writeResponse(response, jsonBody, charset);
                }
            } catch (Exception e) {
                XSUtil.error(this.getClass(), "DWM for [" + apiPath + "] error", e);
            }
        }
    }

    /**
     * 从userRequest中获取apiPath
     * @param request
     * @return
     */
    private String getApiPath(UserRequest request) {
        if (XSUtil.isNotBlank(request.getEndpointApi())) {
            return request.getEndpointApi();
        }
        String apiPath = Baggage.current().getEntryValue("endpoint-api");
        if (XSUtil.isNotBlank(apiPath)) {
            return apiPath;
        }
        Object servletInvocableHandlerMethod = request.getAdviceObject();
        try {
            String methodSign = GenerateSpanAttrUtil.getMethodSign((Class) servletInvocableHandlerMethod.getClass().getMethod("getBeanType").invoke(request.getAdviceObject()), (Method) servletInvocableHandlerMethod.getClass().getMethod("getMethod").invoke(servletInvocableHandlerMethod));
            Api api = ApiCache.cache.get(methodSign);
            StringJoiner stringJoiner = new StringJoiner(",");
            api.getPaths().forEach(path -> stringJoiner.add(path.getName()));
            return stringJoiner.toString();
        } catch (Exception e) {
            XSUtil.error(this.getClass(), e);
        }
        return null;
    }

    /**
     * 调用ApiService数据水印接口获取加水印结果
     * @param url
     * @param jsonBody  目前只支持jsonArray格式
     * @return 返回加水印数据
     */
    private String getDataWaterMarkResult(String url, String jsonBody) {
        XSUtil.debug(this.getClass(), "DWM url:[{}], request：{}", url, jsonBody);
        HttpResponse execute = HttpUtil.createPost(url).body(jsonBody).setConnectionTimeout(3*1000).setReadTimeout(10*1000).execute();
        XSUtil.debug(this.getClass(), "DWM response：{}", execute.body());
        if (execute.isOk()) {
            if ("1001".equals(JsonUtil.getStringByKey(execute.body(), "status"))) {
                return JsonUtil.getStringByKey(execute.body(), "data");
            } else {
                throw new RuntimeException(StrUtil.format("call DWM[{}] fail:{}", url, execute.body()));
            }
        } else {
            throw new RuntimeException(StrUtil.format("call DWM[{}] error:{}", url, execute.toString()));
        }
    }

    /**
     * 将脱敏/水印结果写入到response
     * @param response
     * @param json
     * @param charset
     * @throws IOException
     */
    private void writeResponse(XSHttpServletResponseWrapper response, String json, String charset) throws IOException {
        if (response.getByteArray().length > 0) {
            byte[] bytes = charset != null ? json.getBytes(charset) : json.getBytes();
            ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(bytes.length);
            byteArrayOutputStream.write(bytes);
            response.setByteArrayOutputStream(byteArrayOutputStream);
            response.setContentLength(byteArrayOutputStream.size());
        } else if (response.getCharArray().length > 0) {
            CharArrayWriter charArrayWriter = new CharArrayWriter(json.length());
            charArrayWriter.write(json);
            response.setCharArrayWriter(charArrayWriter);
            response.setContentLength(charArrayWriter.size());
        }
    }

    /**
     * 根据json路径增加对应节点下的数据水印
     * @param url
     * @param jsonBody
     * @param threshold
     * @return
     */
    private String addWaterMarkByPath(String url, String jsonBody, int threshold) {
        String waterMarkPaths = AgentConfigDTO.getConfig().getWaterMarkPaths();
        if (XSUtil.isBlank(waterMarkPaths)) {
            return jsonBody;
        }
        try (JsonReader reader = Json.createReader(new StringReader(jsonBody))) {
            JsonObject jsonObject = reader.readObject();
            for (String path : waterMarkPaths.split(",")) {
                path = path.trim();
                if (XSUtil.isEmpty(path) || !path.startsWith("/")) {
                    continue;
                }
                JsonValue jsonValue = jsonObject.getValue(path);
                if (JsonValue.ValueType.ARRAY != jsonValue.getValueType() || "[]".equals(jsonValue.toString())) {
                    continue;
                }
                if (JsonUtil.getJsonArraySize(jsonValue.toString()) < threshold) {
                    continue;
                }
                String result = getDataWaterMarkResult(url, jsonValue.toString());
                try (JsonReader resultReader = Json.createReader(new StringReader(result))) {
                    JsonArray jsonArray = resultReader.readArray();
                    jsonObject = Json.createPatchBuilder().replace(path, jsonArray).build().apply(jsonObject);
                    jsonBody = jsonObject.toString();
                    XSUtil.debug(this.getClass(), "replace DWM:{}", jsonBody);
                    break;
                }
            }
        }
        return jsonBody;
    }
}
